package com.mrx.tunnel.utils;

import lombok.SneakyThrows;

import java.io.ByteArrayOutputStream;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;

/**
 * @author Mr.X
 * @since 2024/9/2 07:54
 */
public class TunnelUtils {

    private static final int versionNum = 203;

    public static byte[] concatByte(byte[] header, byte[] body) {
        byte[] combined = new byte[header.length + body.length];
        System.arraycopy(header, 0, combined, 0, header.length);
        System.arraycopy(body, 0, combined, header.length, body.length);
        return combined;
    }

    public static byte[] getBlock(String val) {
        byte[] valBytes = val.getBytes(StandardCharsets.UTF_8);
        int len = valBytes.length;
        ByteBuffer buffer;
        if (len < 254) {
            buffer = ByteBuffer.allocate(1 + len);
            buffer.put((byte) len);  // 将长度作为一个字节放入
        } else {
            buffer = ByteBuffer.allocate(1 + 4 + len);
            buffer.put((byte) 0xFE);  // 如果长度 >= 254，首先添加 0xFE
            buffer.putInt(len);       // 然后添加长度的 4 字节表示
        }
        buffer.put(valBytes);  // 添加实际的字符串内容
        return buffer.array();
    }

    public static byte[] getResponse(int errno, String message) {
        byte[] header = getEchoHeader(errno);
        return TunnelUtils.concatByte(header, getBlock(message));
    }

    /**
     * 返回一个查询结果的 header
     *
     * @param errno      错误代码
     * @param affectRows 影响行数
     * @param insertId   插入 id
     * @param fieldNums  字段个数
     * @param rowNums    查询结果行数
     * @return 查询结果 header
     */
    @SneakyThrows
    public static byte[] getEchoResultSetHeader(int errno, long affectRows, long insertId, int fieldNums, int rowNums) {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        outputStream.write(getLongBinary(errno));
        outputStream.write(getLongBinary(affectRows));
        outputStream.write(getLongBinary(insertId));
        outputStream.write(getLongBinary(fieldNums));
        outputStream.write(getLongBinary(rowNums));
        outputStream.write(getDummy(12));
        return outputStream.toByteArray();
    }

    @SneakyThrows
    public static byte[] getEchoHeader(int errno) {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        outputStream.write(getLongBinary(1111));           // 写入固定值 1111
        outputStream.write(getShortBinary(versionNum));    // 写入版本号
        outputStream.write(getLongBinary(errno));          // 写入错误码
        outputStream.write(getDummy(6));                   // 写入 6 个字节的填充数据
        return outputStream.toByteArray();
    }

    private static byte[] getLongBinary(int num) {
        ByteBuffer buffer = ByteBuffer.allocate(4);
        buffer.putInt(num);
        return buffer.array();
    }

    private static byte[] getLongBinary(long num) {
        ByteBuffer buffer = ByteBuffer.allocate(4);
        buffer.putLong(num);
        return buffer.array();
    }

    private static byte[] getShortBinary(int num) {
        ByteBuffer buffer = ByteBuffer.allocate(2);
        buffer.putShort((short) num);
        return buffer.array();
    }

    private static byte[] getDummy(int count) {
        return new byte[count];
    }

}
